import base64
import os
import time
import json
import logging
import aiofiles
from common.core.cron import CronBase
from common.utils.cache import RedisUtil
from common.utils.urls import UrlUtil
from common.utils.times import TimeUtil
from common.utils.tools import ToolsUtil
from common.enums.apply import DrawsEnum
from common.enums.wallet import WalletEnum
from common.models.users import UserWalletModel
from common.models.draws import DrawsRecordsModel
from common.cache.queue_cache import QueueCache
from common.chain.draws_server import DrawsDoubaoServer, DrawsDalleServer
from plugins.storage.driver import StorageDriver

logger = logging.getLogger(__name__)


class Command(CronBase):
    @classmethod
    async def run(cls, **kwargs):
        try:
            # 读取任务
            record_id = await QueueCache.queue_pop(QueueCache.DRAW_JOB)
            if not record_id:
                return

            # 查询任务
            record = await DrawsRecordsModel.filter(id=int(record_id)).first()
            if not record or record.status != DrawsEnum.STATUS_ING:
                return

            # 处理响应
            if record.engine == "doubao":
                await cls.handle_doubao_check(record)
            elif record.engine == "dalle":
                await cls.handle_dalle_check(record)
            elif record.engine == "sd":
                pass
            else:
                record.status = DrawsEnum.STATUS_FAIL
                record.fail_reason = f"No processing channel {record.engine}"
                await record.save()
        except Exception as e:
            logger.error(f"handle_draws error: {str(e)}")

    @classmethod
    async def handle_doubao_check(cls, record: DrawsRecordsModel):
        album_dir_path = ""
        thumb_dir_path = ""
        try:
            # 限流拦截
            lock_key: str = f"prevent:draws_{str(record.id)}"
            lock_ack: bool = await RedisUtil.set(lock_key, str(record.id), 20, True)
            if not lock_ack:
                await QueueCache.queue_push(QueueCache.DRAW_JOB, record.id)
                return False

            # 获取生成结果
            server = DrawsDoubaoServer()
            response = await server.fetch(record.model, record.task_id)

            # 根据状态处理
            if response.get("status") == "done":
                image_url = response.get("image_urls")[0]
                await DrawsRecordsModel.filter(id=record.id).update(
                    status=DrawsEnum.STATUS_YES,
                    image_url=image_url,
                    thumbnail=image_url,
                    notify_snap=json.dumps(response, ensure_ascii=False),
                    update_time=int(time.time())
                )

                # 图片扩展
                urlName: str = os.path.basename(image_url.split("?")[0])
                fileExt: str = str(urlName.split(".")[-1]).strip().lower()
                if fileExt not in ["png", "jpg", "jpeg", "gif", "webp"]:
                    fileExt = "jpg"

                # 存储路径
                filename: str = ToolsUtil.make_md5_str(urlName + str(time.time())) + "." + fileExt
                curr_now: str = TimeUtil.timestamp_to_date(int(time.time()), str("%Y%m%d"))
                album_dir: str = f"storage/applies/draws/albums/{curr_now}/{filename}"
                thumb_dir: str = f"storage/applies/draws/thumbnail/{str(curr_now)}/{filename}"
                album_dir_path: str = UrlUtil.to_root_path(path=album_dir)
                thumb_dir_path: str = UrlUtil.to_root_path(path=thumb_dir)

                # 下载图片
                await ToolsUtil.download_file(url=image_url, save_path=album_dir_path)

                # 生成缩略图
                ToolsUtil.generate_thumbnail(input_path=album_dir_path, output_path=thumb_dir_path)

                # 图上传OSS
                engine: str = await UrlUtil.get_storage_engine()
                if engine != str("local"):
                    await StorageDriver.push(path=album_dir_path, key=album_dir)
                    await StorageDriver.push(path=thumb_dir_path, key=thumb_dir)

                # 更新图片
                await DrawsRecordsModel.filter(id=record.id).update(
                    image_loc=album_dir,
                    thumbnail=thumb_dir
                )
            elif response.get("status") == "in_queue":
                await QueueCache.queue_push(QueueCache.DRAW_JOB, record.id)
            else:
                await DrawsRecordsModel.filter(id=record.id).update(
                    status=DrawsEnum.STATUS_FAIL,
                    fail_reason=response.get("message"),
                    notify_snap=json.dumps(response, ensure_ascii=False),
                    update_time=int(time.time())
                )
                await UserWalletModel.inc(
                    scene=WalletEnum.POINTS,
                    user_id=record.user_id,
                    change_type=WalletEnum.POINTS_INC_DRAWS_FAIL,
                    change_amount=record.use_points,
                    project="AI绘画",
                    remarks=response.get("message")
                )
        except Exception as e:
            if record.retry_count < 2:
                record.retry_count += 1
                record.start_time = int(time.time())
                await record.save()
                await QueueCache.queue_push(QueueCache.DRAW_JOB, record.id)
            else:
                await DrawsRecordsModel.filter(id=record.id).update(
                    fail_reason=str(e),
                    status=DrawsEnum.STATUS_FAIL,
                    notify_snap=json.dumps({}, ensure_ascii=False),
                    update_time=int(time.time())
                )
                if record.use_points > 0:
                    await UserWalletModel.inc(
                        scene=WalletEnum.POINTS,
                        user_id=record.user_id,
                        change_type=WalletEnum.POINTS_INC_DRAWS_FAIL,
                        change_amount=record.use_points,
                        project="AI绘画",
                        remarks=str(e)
                    )
        finally:
            engine: str = await UrlUtil.get_storage_engine()
            if engine != "local":
                if album_dir_path and os.path.exists(album_dir_path):
                    os.unlink(album_dir_path)
                if thumb_dir_path and os.path.exists(thumb_dir_path):
                    os.unlink(thumb_dir_path)

    @classmethod
    async def handle_dalle_check(cls, record: DrawsRecordsModel):
        album_dir_path = ""
        thumb_dir_path = ""
        try:
            # 限流拦截
            lock_key: str = f"prevent:draws_{str(record.id)}"
            lock_ack: bool = await RedisUtil.set(lock_key, str(record.id), 20, True)
            if not lock_ack:
                await QueueCache.queue_push(QueueCache.DRAW_JOB, record.id)
                return False

            # 初始通道
            dalle_server = DrawsDalleServer({
                "model": record.model,
                "prompt": record.prompt_in,
                "scale": record.scale,
                "other_params": json.loads(record.other_params)
            })

            # 发起绘画
            response = await dalle_server.imagine()
            if not response.get("data"):
                error = response.get("error", {})
                raise Exception(error.get("message", "响应数据异常"))

            # 存储路径
            filename: str = ToolsUtil.make_md5_str(record.prompt_in + str(time.time())) + ".png"
            curr_now: str = TimeUtil.timestamp_to_date(int(time.time()), str("%Y%m%d"))
            album_dir: str = f"storage/drawing/albums/{curr_now}/{filename}"
            thumb_dir: str = f"storage/drawing/thumbnail/{str(curr_now)}/{str(filename)}"
            album_dir_path: str = UrlUtil.to_root_path(path=album_dir)
            thumb_dir_path: str = UrlUtil.to_root_path(path=thumb_dir)
            if not os.path.isdir(os.path.dirname(album_dir_path)):
                os.makedirs(os.path.dirname(album_dir_path), exist_ok=True)

            # 原图保存
            b64_json = response["data"][0]["b64_json"]
            b64_data = base64.b64decode(b64_json.split(",")[1])
            async with aiofiles.open(album_dir_path, "wb") as file:
                await file.write(b64_data)

            # 生成缩略图
            ToolsUtil.generate_thumbnail(input_path=album_dir_path, output_path=thumb_dir_path)

            # 图上传OSS
            engine: str = await UrlUtil.get_storage_engine()
            if engine != str("local"):
                await StorageDriver.push(path=album_dir_path, key=album_dir)
                await StorageDriver.push(path=thumb_dir_path, key=thumb_dir)

            # 更新状态
            await DrawsRecordsModel.filter(id=record.id).update(
                status=DrawsEnum.STATUS_YES,
                image_loc=album_dir,
                image_url=album_dir,
                thumbnail=thumb_dir,
                notify_snap=json.dumps({}, ensure_ascii=False),
                update_time=int(time.time())
            )
        except Exception as e:
            if record.retry_count < 2:
                record.retry_count += 1
                record.start_time = int(time.time())
                await record.save()
                await QueueCache.queue_push(QueueCache.DRAW_JOB, record.id)
            else:
                await DrawsRecordsModel.filter(id=record.id).update(
                    status=DrawsEnum.STATUS_FAIL,
                    fail_reason=str(e),
                    notify_snap=json.dumps({}, ensure_ascii=False),
                    update_time=int(time.time())
                )
                if record.use_points > 0:
                    await UserWalletModel.inc(
                        scene=WalletEnum.POINTS,
                        user_id=record.user_id,
                        change_type=WalletEnum.POINTS_INC_DRAWS_FAIL,
                        change_amount=record.use_points,
                        project=str("AI绘画"),
                        remarks=str(e)
                    )
        finally:
            engine: str = await UrlUtil.get_storage_engine()
            if engine != "local":
                if album_dir_path and os.path.exists(album_dir_path):
                    os.unlink(album_dir_path)
                if thumb_dir_path and os.path.exists(thumb_dir_path):
                    os.unlink(thumb_dir_path)
